from slip_detector import SlipDeltaMeter
import numpy as np
import os
import matplotlib.pyplot as plt
from collections import deque

data_dir = r"datalib/dataset_1"

def load_data(exp_dir, sensor_id):
    try:
        # print(f"Loading data from {exp_dir} for sensor {sensor_id}")
        tacdata = np.load(os.path.join(exp_dir, f"tacdata_{sensor_id}.npy"))
        true_delta = np.load(os.path.join(exp_dir, f"true_delta_{sensor_id}.npy"))
        predicted_delta = np.load(os.path.join(exp_dir, f"predicted_delta_{sensor_id}.npy"))
        grasp_force = np.load(os.path.join(exp_dir, f"grasp_force_{sensor_id}.npy"))
    
        return tacdata, true_delta, predicted_delta, grasp_force
    except FileNotFoundError as e:
        print(f"File not found: {e}")
        return None, None, None, None

def calc_contact_center(tacdata):
    """
    Calculate the contact center from Tac3D data.
    
    Args:
        tacdata (np.ndarray): The Tac3D data array.
        
    Returns:
        np.ndarray: The contact center coordinates.
    """
    F = np.linalg.norm(tacdata[:, 6:9], axis=1)
    contact_mask = F > 0.01
    contact_mask = contact_mask.reshape(-1,20,20)
    contact_positions = np.stack(np.where(contact_mask))
    return np.mean(contact_positions, axis=0)

class TangentForceRateCalculator:
    """Calculate the rate of change of tangent force."""
    def __init__(self, window_size=5, momentum=0.9):
        """
        Initialize the calculator with a window size.
        Args:
            window_size (int): The size of the window for calculating the rate.
        """
        self.window_size = window_size
        self.previous_tangent_force = deque(maxlen=window_size)
        self.momentum = momentum
        self.now_rate = 0.0
        x = np.arange(window_size)
        self.A = np.vstack([x, np.ones(len(x))]).T
    
    def reset(self):
        """
        Reset the calculator.
        """
        self.previous_tangent_force.clear()
        self.now_rate = 0.0
    
    def calculate_rate(self, tangent_force : float):
        """
        Calculate the rate of change of tangent force.
        Args:
            tangent_force (float): The current tangent force value.
        Returns:
            float: The rate of change of tangent force.
        """

        self.previous_tangent_force.append(tangent_force)
        if len(self.previous_tangent_force) < self.window_size:
            return 0.0
        
        # Calculate the rate of change with all the data in the window using linear fitting

        y = np.array(self.previous_tangent_force)

        m, c = np.linalg.lstsq(self.A, y, rcond=None)[0]
        self.now_rate = self.now_rate * self.momentum + m * (1 - self.momentum)
        return self.now_rate
        

exp_dir_list = []
for dir in os.listdir(data_dir):
    exp_dir = os.path.join(data_dir, dir)
    if not os.path.isdir(exp_dir):
        continue
    exp_dir_list.append(exp_dir)

# sort the directories by their names
exp_dir_list.sort(key=lambda x: int(os.path.basename(x).split('_')[1]), reverse=True)
rate_calculator = TangentForceRateCalculator(window_size=5, momentum=0.9)
for exp_dir in exp_dir_list:

    plt.figure(figsize=(10, 9))
    for sensor_id in [0,1]:
        tacdata, true_delta, predicted_delta, grasp_force = load_data(exp_dir, sensor_id)
        if tacdata is None:
            continue
        rate_calculator.reset()
        sensored_force = np.sum(tacdata[:, :, 6:9], axis=1)
        tangent_force = np.linalg.norm(sensored_force[:,0:2], axis=1)
        normal_force = -sensored_force[:, 2]

        tangent_force_rate = []
        for tf in tangent_force:
            rate = rate_calculator.calculate_rate(tf)
            tangent_force_rate.append(rate)

        # contact_posx = []
        # contact_posy = []
        # for i in range(tacdata.shape[0]):
        #     contact_center = calc_contact_center(tacdata[i])
        #     contact_posx.append(contact_center[0])
        #     contact_posy.append(contact_center[1])


        plt.subplot(2, 2, sensor_id + 1)
        plt.plot(tangent_force, label='Tangent Force')

        plt.plot(normal_force, label='Normal Force')
        plt.plot(grasp_force, label='Grasp Force')
        plt.ylabel('Force (N)')
        plt.legend()

        plt.subplot(2, 2, sensor_id + 3)
        plt.plot(tangent_force_rate, label='Tangent Force Rate', linestyle='--')
        # plt.plot(contact_posx, label='Contact Center X')
        # plt.plot(contact_posy, label='Contact Center Y')
        # plt.plot(true_delta, label='True Delta', color='blue', linewidth=2)
        # plt.plot(predicted_delta, label='Estimated Delta', color='orange', linestyle='--', linewidth=2)
        # plt.title('Slip Delta Estimation')
        plt.xlabel('Sample Index')
        # plt.ylabel('Slip Delta')
        plt.legend()
        mae = np.mean(np.abs(predicted_delta - true_delta))
        print(f"Sensor {sensor_id} MAE: {mae:.4f}")
    plt.show()



# tacdata = np.load(os.path.join(data_dir, "tacdata_1.npy"))
# true_delta = np.load(os.path.join(data_dir, "true_delta_1.npy"))

# grasp_force = np.load(os.path.join(data_dir, "grasp_force_1.npy"))
# sensored_force = np.sum(tacdata[:, :, 6:9], axis=1)
# tangent_force = np.linalg.norm(sensored_force[:,0:2], axis=1)
# normal_force = -sensored_force[:, 2]

# predicted_delta = np.load(os.path.join(data_dir, "predicted_delta_1.npy"))

# plt.subplot(2, 2, 2)
# plt.plot(tangent_force, label='Tangent Force')
# plt.plot(normal_force, label='Normal Force')
# plt.plot(grasp_force, label='Grasp Force')
# # plt.plot(grad_grasp_force, label='Gradient of Grasp Force')
# plt.ylabel('Force (N)')
# plt.legend()

# plt.subplot(2, 2, 4)
# plt.plot(true_delta, label='True Delta', color='blue', linewidth=2)
# plt.plot(predicted_delta, label='Estimated Delta', color='orange', linestyle='--', linewidth=2)
# plt.title('Slip Delta Estimation')
# plt.xlabel('Sample Index')
# plt.ylabel('Slip Delta')
# plt.legend()

# plt.show()